import java.lang.foreign.Arena;
import java.lang.foreign.FunctionDescriptor;
import java.lang.foreign.Linker;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.SymbolLookup;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.lang.reflect.AccessibleObject;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Arrays;
import static java.lang.foreign.ValueLayout.*;

//ref: https://github.com/dreamlike-ocean/backend_qingyou/blob/main/dreamlike%E7%9A%84%E7%A7%81%E8%B4%A7/afterUnsafe.md
//ref: https://github.com/dreamlike-ocean/UnsafeJava/blob/master/unsafe-core/src/main/java/top/dreamlike/unsafe/core/panama/jni/JNIEnv.java
public class Jvm {
	public static final MethodHandles.Lookup lookup = MethodHandles.lookup();

	static {
		var jvmName = System.mapLibraryName("jvm");
		var javaHome = System.getProperty("java.home");
		var jvmPath = javaHome + "/lib/server/" + jvmName;
		if (!Files.exists(Path.of(jvmPath)))
			jvmPath = javaHome + "/bin/server/" + jvmName;
		Runtime.getRuntime().load(jvmPath);

		try (var arena = Arena.ofConfined()) {
			var pAddr = arena.allocate(ADDRESS);
			var pNumVms = arena.allocate(JAVA_INT);
			pNumVms.set(JAVA_INT, 0, 0);
			var linker = Linker.nativeLinker();
			// jint JNI_GetCreatedJavaVMs(JavaVM** vm_buf, jsize bufLen, jsize* numVMs);
			int r = (int)linker.downcallHandle(SymbolLookup.loaderLookup().find("JNI_GetCreatedJavaVMs").orElseThrow(),
					FunctionDescriptor.of(JAVA_INT, ADDRESS, JAVA_INT, ADDRESS)).invokeExact(pAddr, 1, pNumVms);
			if (r != 0)
				throw new ExceptionInInitializerError("JNI_GetCreatedJavaVMs=" + r);
			var pJavaVm = pAddr.get(ADDRESS, 0);
			var pGetEnv = pJavaVm.reinterpret(ADDRESS.byteSize()).get(ADDRESS, 0).reinterpret(ADDRESS.byteSize() * 8)
					.get(ADDRESS, ADDRESS.byteSize() * 6);
			// int GetEnv(JavaVM* vm, void** penv, jint version);
			r = (int)linker.downcallHandle(pGetEnv, FunctionDescriptor.of(JAVA_INT, ADDRESS, ADDRESS, JAVA_INT))
					.invokeExact(pJavaVm, pAddr, 21 << 16);
			if (r != 0)
				throw new ExceptionInInitializerError("JavaVM.GetEnv=" + r);
			var pEnv = pAddr.get(ADDRESS, 0);
			var envFuncs = pEnv.reinterpret(ADDRESS.byteSize()).get(ADDRESS, 0).reinterpret(ADDRESS.byteSize() * 235);

			// jclass FindClass(JNIEnv* env, const char* name);
			var mhFindClass = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 6),
					FunctionDescriptor.of(ADDRESS, ADDRESS, ADDRESS));
			// jobject NewGlobalRef(JNIEnv* env, jobject lobj);
			var mhNewGlobalRef = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 21),
					FunctionDescriptor.of(ADDRESS, ADDRESS, ADDRESS));
			// void DeleteGlobalRef(JNIEnv* env, jobject gref);
			var mhDeleteGlobalRef = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 22),
					FunctionDescriptor.ofVoid(ADDRESS, ADDRESS));
			// jmethodID GetMethodID(JNIEnv* env, jclass clazz, const char* name, const char* sig);
			var mhGetMethodID = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 33),
					FunctionDescriptor.of(ADDRESS, ADDRESS, ADDRESS, ADDRESS, ADDRESS));
			// jobject CallObjectMethod(JNIEnv* env, jobject obj, jmethodID methodID, ...);
			var mhCallObjectMethod0 = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 34),
					FunctionDescriptor.of(ADDRESS, ADDRESS, ADDRESS, ADDRESS));
			// jmethodID GetStaticMethodID(JNIEnv* env, jclass clazz, const char* name, const char* sig);
			var mhGetStaticMethodID = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 113),
					FunctionDescriptor.of(ADDRESS, ADDRESS, ADDRESS, ADDRESS, ADDRESS));
			// jobject CallStaticObjectMethod(JNIEnv* env, jclass clazz, jmethodID methodID, ...);
			var mhCallStaticObjectMethod0 = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 114),
					FunctionDescriptor.of(ADDRESS, ADDRESS, ADDRESS, ADDRESS));
			var mhCallStaticObjectMethod3 = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 114),
					FunctionDescriptor.of(ADDRESS, ADDRESS, ADDRESS, ADDRESS, ADDRESS, ADDRESS, ADDRESS));
			// jfieldID GetStaticFieldID(JNIEnv* env, jclass clazz, const char* name, const char* sig);
			var mhGetStaticFieldID = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 144),
					FunctionDescriptor.of(ADDRESS, ADDRESS, ADDRESS, ADDRESS, ADDRESS));
			// jobject GetStaticObjectField(JNIEnv* env, jclass clazz, jfieldID fieldID);
			var mhGetStaticObjectField = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 145),
					FunctionDescriptor.of(ADDRESS, ADDRESS, ADDRESS, ADDRESS));
			// void SetStaticObjectField(JNIEnv* env, jclass clazz, jfieldID fieldID, jobject value);
			var mhSetStaticObjectField = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 154),
					FunctionDescriptor.ofVoid(ADDRESS, ADDRESS, ADDRESS, ADDRESS));
			// jstring NewStringUTF(JNIEnv* env, const char* utf);
			var mhNewStringUTF = linker.downcallHandle(envFuncs.get(ADDRESS, ADDRESS.byteSize() * 167),
					FunctionDescriptor.of(ADDRESS, ADDRESS, ADDRESS));

			var threadClass = (MemorySegment)mhFindClass.invokeExact(pEnv,
					arena.allocateFrom("java/lang/Thread"));
			var threadClassGlobal = (MemorySegment)mhNewGlobalRef.invokeExact(pEnv, threadClass);
			var classClass = (MemorySegment)mhFindClass.invokeExact(pEnv,
					arena.allocateFrom("java/lang/Class"));
			var classClassGlobal = (MemorySegment)mhNewGlobalRef.invokeExact(pEnv, classClass);
			var lookupClass = (MemorySegment)mhFindClass.invokeExact(pEnv,
					arena.allocateFrom("java/lang/invoke/MethodHandles$Lookup"));
			var lookupClassGlobal = (MemorySegment)mhNewGlobalRef.invokeExact(pEnv, lookupClass);

			var currentThreadMethodId = (MemorySegment)mhGetStaticMethodID.invokeExact(pEnv, threadClassGlobal,
					arena.allocateFrom("currentThread"), arena.allocateFrom("()Ljava/lang/Thread;"));
			var getContextClassLoaderMethodId = (MemorySegment)mhGetMethodID.invokeExact(pEnv, threadClassGlobal,
					arena.allocateFrom("getContextClassLoader"), arena.allocateFrom("()Ljava/lang/ClassLoader;"));
			var forNameMethodId = (MemorySegment)mhGetStaticMethodID.invokeExact(pEnv, classClassGlobal,
					arena.allocateFrom("forName"),
					arena.allocateFrom("(Ljava/lang/String;ZLjava/lang/ClassLoader;)Ljava/lang/Class;"));

			var curThreadObj = (MemorySegment)mhCallStaticObjectMethod0.invokeExact(
					pEnv, threadClassGlobal, currentThreadMethodId);
			var curThreadObjGlobal = (MemorySegment)mhNewGlobalRef.invokeExact(pEnv, curThreadObj);
			var curClassLoaderObj = (MemorySegment)mhCallObjectMethod0.invokeExact(
					pEnv, curThreadObjGlobal, getContextClassLoaderMethodId);
			var curClassLoaderObjGlobal = (MemorySegment)mhNewGlobalRef.invokeExact(pEnv, curClassLoaderObj);

			var thisClassNameStr = (MemorySegment)mhNewStringUTF.invokeExact(
					pEnv, arena.allocateFrom(Jvm.class.getName().replace('.', '/')));
			var thisClassNameStrGlobal = (MemorySegment)mhNewGlobalRef.invokeExact(pEnv, thisClassNameStr);
			var thisClassObj = (MemorySegment)mhCallStaticObjectMethod3.invokeExact(pEnv, classClassGlobal,
					forNameMethodId, thisClassNameStrGlobal, MemorySegment.NULL, curClassLoaderObjGlobal);
			var thisClassObjGlobal = (MemorySegment)mhNewGlobalRef.invokeExact(pEnv, thisClassObj);

			var lookupSig = arena.allocateFrom("Ljava/lang/invoke/MethodHandles$Lookup;");
			var implLookupFieldId = (MemorySegment)mhGetStaticFieldID
					.invokeExact(pEnv, lookupClassGlobal, arena.allocateFrom("IMPL_LOOKUP"), lookupSig);
			var thisLookupFieldId = (MemorySegment)mhGetStaticFieldID
					.invokeExact(pEnv, thisClassObjGlobal, arena.allocateFrom("lookup"), lookupSig);

			mhSetStaticObjectField.invokeExact(pEnv, thisClassObjGlobal, thisLookupFieldId,
					(MemorySegment)mhGetStaticObjectField.invokeExact(pEnv, lookupClassGlobal, implLookupFieldId));

			mhDeleteGlobalRef.invokeExact(pEnv, thisClassObjGlobal);
			mhDeleteGlobalRef.invokeExact(pEnv, thisClassNameStrGlobal);
			mhDeleteGlobalRef.invokeExact(pEnv, curClassLoaderObjGlobal);
			mhDeleteGlobalRef.invokeExact(pEnv, curThreadObjGlobal);
			mhDeleteGlobalRef.invokeExact(pEnv, lookupClassGlobal);
			mhDeleteGlobalRef.invokeExact(pEnv, classClassGlobal);
			mhDeleteGlobalRef.invokeExact(pEnv, threadClassGlobal);
		} catch (RuntimeException | Error e) {
			throw e;
		} catch (Throwable e) {
			throw new ExceptionInInitializerError(e);
		}
	}

	public static void main(String[] args) throws Throwable {
		System.out.println(System.getProperty("java.home"));
		System.out.println(lookup); // "/trusted"

		var unsafeClass = Class.forName("jdk.internal.misc.Unsafe");
		var u = lookup.findStaticGetter(unsafeClass, "theUnsafe", unsafeClass).invoke();
		System.out.println(u);

		var getFields = lookup.findSpecial(Class.class, "getDeclaredFields0",
				MethodType.methodType(Field[].class, boolean.class), Class.class);
		var getMethods = lookup.findSpecial(Class.class, "getDeclaredMethods0",
				MethodType.methodType(Method[].class, boolean.class), Class.class);
		System.out.println("---");
		System.out.println(Arrays.toString((Field[])getFields.invokeExact(
				AccessibleObject.class, false)).replace(", ", "\n "));
		System.out.println("---");
		System.out.println(Arrays.toString((Method[])getMethods.invokeExact(
				Class.forName("jdk.internal.reflect.Reflection"), false)).replace(", ", "\n "));
		System.out.println("---");
	}
}
